# 从csv文件中读取数据，分别为：X列表和对应的Y列表

import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression


def get_data():
    # 加载CSV文件
    df = pd.read_csv('./linear_regression/auto_mpg.csv')

    # 提取特征列和标签列（注意题目要求：用重量预测MPG）
    X = df[['weight']]  # 特征向量（车辆重量kg）DataFrame格式
    Y = df['mpg']  # 目标向量（每加仑英里数）

    return X, Y


def get_linear_model(X, Y):
    # 创建线性回归模型
    model = LinearRegression()

    # 训练模型
    model.fit(X, Y)

    # 输出方程参数
    print(f"回归方程：mpg = {model.coef_[0]:.4f} * weight + {model.intercept_:.2f}")

    return model


def show_linear_line(X, Y, model):
    plt.figure(figsize=(10, 6))

    # 1. 绘制原始数据散点图
    plt.scatter(X, Y, color='blue', alpha=0.5, label='Actual Data')

    # 2. 绘制预测直线
    predictions = model.predict(X)
    plt.plot(X, predictions, color='red', linewidth=2, label='Regression Line')

    # 图表装饰
    plt.title('Vehicle MPG Prediction by Weight')
    plt.xlabel('Weight (kg)')
    plt.ylabel('Miles Per Gallon (MPG)')
    plt.legend()
    plt.grid(True)

    plt.show()


def main():
    # 1. 读取数据
    X, Y = get_data()

    # 2. 模型训练
    model = get_linear_model(X, Y)

    # 3. 获取车辆重量为3000时的预测行驶距离
    weight_3000 = pd.DataFrame([[3000]], columns=['weight'])  # 指定列名
    predicted_mpg = model.predict(weight_3000)
    print(f'\n车辆重量3000kg时的预测MPG：{predicted_mpg[0]:.2f} 英里/加仑')

    # 4. 绘图
    show_linear_line(X, Y, model)



if __name__ == '__main__':
    main()